from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from utils.tools import read_json
from utils.consts import DATA_CONF, DATA_PATH

r'''
assert img data is placed as format:
    root
    ├── conf.json
    ├── train
    │   ├── airplane
    │   ├── ...
    │   └── truck
    └── val
        ├── airplane
        ├── ...
        └── truck

distr = [1, 0.5, 0.3333, 0.25, 0.2, 0, 0, 0, 0, 0]
assert len(distr) = len(cls_dict)
'''


class AdjustSet(Dataset):
    def __init__(self, root: str, distr: list, transform, train) -> None:
        super().__init__()
        # path -- example : '../data/stl_10' -> '../data/stl_10/train'
        self.img_path = f"{root}/train" if train else f"{root}/val"

        r'''
        Origin datainfo, example:
            cls_dict = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
            org_distr = [5040, 5040, 5040, 5040, 5040, 5040, 5040, 5040, 5040, 5040]
            org_num_cls = 10
            distr = [5040, 2520, 1680, 1260, 1008, 0, 0, 0, 0, 0]
        '''
        data_conf = read_json(f"{root}/{DATA_CONF}")
        cls_dict = data_conf['label']
        org_distr = data_conf['train'] if train else data_conf['val']
        org_num_cls = len(cls_dict)
        if train:
            distr = [round(distr[i]*org_distr[i]) for i in range(org_num_cls)]

        r'''
        Set distr dataset, example:
            max_distr = [5040, 5040, 5040, 5040, 5040]
            self.distr = [5040, 2520, 1680, 1260, 1008]
            self.cls_dict = ["airplane", "automobile", "bird", "cat", "deer"]
            self.num_cls = 5
        '''
        self.ask_distr = distr
        max_distr = [org_distr[i] for i in range(org_num_cls) if distr[i] != 0]
        self.distr = [x for x in distr if x != 0] if train else max_distr
        self.cls_dict = [cls_dict[i]
                         for i in range(org_num_cls) if distr[i] != 0]
        self.num_cls = len(self.cls_dict)
        self.data = self.get_imgnames(max_distr)
        self.transform = transform

    # get ids = [0, ... ,maxid-1], use rate.5 to get window[left, right) to slice
    # so that we can get data which satisfied distribution
    def get_imgnames(self, max_distr):
        data = []
        for cls_id in range(self.num_cls):
            max_num, ask_num = max_distr[cls_id], self.distr[cls_id]
            id_list = [x for x in range(max_num)]
            left, right = AdjustSet.get_window(max_num-ask_num, ask_num)
            id_list = id_list[left:right]
            for id in id_list:
                data.append((id, cls_id))
        return data

    def get_window(diff: int, len: int, rate: float = 0.5) -> tuple[int, int]:
        left = round(diff * rate)
        right = left + len
        return left, right

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        pic_id, cls_id = self.data[index]
        pic = Image.open(
            f"{self.img_path}/{self.cls_dict[cls_id]}/{pic_id}.png")

        if self.transform != None:
            pic = self.transform(pic)

        return pic, cls_id


def get_loaders(name: str, distr: list, batch_size: int, input_size: int):
    tr_loader = get_loader(name, distr, batch_size, input_size, True)
    vl_loader = get_loader(name, distr, batch_size, input_size, False)
    return tr_loader, vl_loader


def get_loader(name: str, distr: list, batch_size: int, input_size: int, train=True):
    dataset = get_dataset(name, distr, input_size, train)
    return DataLoader(dataset, batch_size, shuffle=train)


def get_datasets(name: str, distr: list, input_size: int):
    tr_set = get_dataset(name, distr, input_size=input_size, train=True)
    vl_set = get_dataset(name, distr, input_size=input_size, train=False)
    return tr_set, vl_set


def get_dataset(name: str, distr: list, input_size: int, train=True):
    transform = transforms.Compose([
        transforms.Resize((input_size, input_size)),
        transforms.ToTensor()
    ])
    root = f"{DATA_PATH}/{name}"
    return AdjustSet(root, distr, transform, train)
